package com.ruoyi.common.utils;

import cn.hutool.core.collection.CollectionUtil;
import com.google.gson.JsonObject;
import com.ruoyi.common.core.domain.entity.MilvusSchema;
import com.ruoyi.common.utils.spring.SpringUtils;
import io.milvus.common.clientenum.FunctionType;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.ConsistencyLevel;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.*;
import io.milvus.v2.service.partition.request.CreatePartitionReq;
import io.milvus.v2.service.partition.request.DropPartitionReq;
import io.milvus.v2.service.partition.request.HasPartitionReq;
import io.milvus.v2.service.partition.request.ReleasePartitionsReq;
import io.milvus.v2.service.vector.request.*;
import io.milvus.v2.service.vector.request.data.BaseVector;
import io.milvus.v2.service.vector.request.data.EmbeddedText;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.request.data.SparseFloatVec;
import io.milvus.v2.service.vector.request.ranker.BaseRanker;
import io.milvus.v2.service.vector.request.ranker.WeightedRanker;
import io.milvus.v2.service.vector.response.DeleteResp;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.SearchResp;
import io.milvus.v2.service.vector.response.UpsertResp;

import java.util.*;

/**
 * 使用milvus
 * 1.如果有相同架构的实体数据就采用同一个集合做不同分区来进行数据隔离（milvus 推荐使用一个集合来存储数据）
 * 2.减少milvus 字段定义，一般只保留必须过滤字段（目前这一版只有一个dm25稀疏向量和lvf_float稠密向量如需更多请修改注释）
 * 3.milvus 无返回值的方法有一个特性 失败时候会抛出异常 所以我们采用关系型数据库事务的方式来解决milvus与pgsql 数据 同步
 */

public class MilvusUtils {

    public static MilvusClientV2 getMilvusClientV2() {
        return SpringUtils.getBean(MilvusClientV2.class);
    }


    /**
     * 创建集合 （先创建CreateCollectionReq.CollectionSchema 然后创建IndexParam）
     *
     * @param collectionName   集合名称
     * @param collectionSchema 字段
     * @param indexParam       字段索引
     */
    public static void createCollection(String collectionName, CreateCollectionReq.CollectionSchema collectionSchema, List<IndexParam> indexParam) {

        if (!hasCollection(collectionName)) {
            CreateCollectionReq build = CreateCollectionReq.builder()
                    .collectionName(collectionName)
                    .collectionSchema(collectionSchema)
                    .indexParams(indexParam)
                    .build();
            getMilvusClientV2().createCollection(build);
        }
    }

    /**
     * 加载集合
     *
     * @param collectionName 集合名称
     */
    public static void loadCollection(String collectionName) {
        getMilvusClientV2().loadCollection(LoadCollectionReq
                .builder()
                .collectionName(collectionName)
                .build());

    }

    /**
     * 判断集合存不存在
     *
     * @param collectionName 集合名
     * @return boolean
     */

    public static boolean hasCollection(String collectionName) {
        HasCollectionReq hasCollectionReq = HasCollectionReq.builder()
                .collectionName(collectionName)
                .build();
        return getMilvusClientV2().hasCollection(hasCollectionReq);
    }

    /**
     * 创建分区
     *
     * @param collectionName 集合名
     * @param partitionName  分区名
     */
    public static void createPartition(String collectionName, String partitionName) {

        //分区不存在就创建
        if (!hasPartition(collectionName, partitionName)) {
            CreatePartitionReq createPartitionReq = CreatePartitionReq.builder()
                    .collectionName(collectionName)
                    .partitionName(partitionName)
                    .build();
            getMilvusClientV2().createPartition(createPartitionReq);
        }

    }

    /**
     * 释放分区（想要删除分区或者停用分区必须要释放分区）
     *
     * @param collectionName 集合名
     * @param partitionName  分区名
     */
    public static void releasePartition(String collectionName, String partitionName) {
        //分区不存在就创建
        if (!hasPartition(collectionName, partitionName)) {
            ReleasePartitionsReq releasePartitionsReq = ReleasePartitionsReq.builder()
                    .collectionName(collectionName)
                    .partitionNames(Collections.singletonList(partitionName))
                    .build();
            getMilvusClientV2().releasePartitions(releasePartitionsReq);
        }
    }

    /**
     * 删除分区
     *
     * @param collectionName 集合名称
     * @param partitionName  分区名称
     * @throws InterruptedException 线程异常
     */
    public static void deletePartition(String collectionName, String partitionName) throws InterruptedException {
        if (hasPartition(collectionName, partitionName)) {
            //删除分区需要先释放分区数据
            getMilvusClientV2().releasePartitions(ReleasePartitionsReq
                    .builder()
                    .collectionName(collectionName)
                    .partitionNames(Collections.singletonList(partitionName))
                    .build()
            );
            //判断集合是否被释放
            boolean state = true;
            while (state) {
                state = getMilvusClientV2().getLoadState(GetLoadStateReq.builder()
                        .collectionName(collectionName)
                        .partitionName(partitionName)
                        .build()
                );
                Thread.sleep(1000);
            }
            //删除分区
            getMilvusClientV2().dropPartition(DropPartitionReq.builder()
                    .partitionName(partitionName)
                    .collectionName(collectionName)
                    .build()
            );

        }
    }

    /**
     * 判断分区是否存在
     *
     * @param collectionName 集合名
     * @param partitionName  分区名
     * @return boolean
     */
    public static boolean hasPartition(String collectionName, String partitionName) {
        HasPartitionReq hasPartitionReq = HasPartitionReq.builder()
                .collectionName(collectionName)
                .partitionName(partitionName)
                .build();
        return getMilvusClientV2().hasPartition(hasPartitionReq);
    }


    /**
     * 创建集合字段请求
     *
     * @param schemas 所有字段
     * @return 集合创建请求
     */
    public static CreateCollectionReq.CollectionSchema createSchema(List<MilvusSchema> schemas) {
        List<AddFieldReq> addFieldReqs = new ArrayList<>();
        //数据库中定义的字段类型进行转换
        schemas.forEach(schema -> {
            //创建字段构建器
            AddFieldReq.AddFieldReqBuilder<?, ?> builder = AddFieldReq.builder();
            builder.fieldName(schema.getFieldName());
            builder.description(schema.getDescription());
            builder.dataType(DataType.forNumber(schema.getDataType()));
            //当设置了字段长度时候给字段长度
            if (schema.getMaxLength() > 0) {
                builder.maxLength(schema.getMaxLength());
            }

            //当字段类型是FloatVector向量类型时候就给维度
            if (schema.getDataType() == DataType.FloatVector.getCode()) {
                builder.dimension(schema.getDimension());
            }
            builder.isPrimaryKey(schema.isPrimaryKey());
            builder.isPartitionKey(schema.isPartitionKey());
            builder.autoID(schema.isAutoID());

            //当数据类型是varchar 时候开启内部分析器
            if (schema.getDataType() == DataType.VarChar.getCode()) {
                //定义中文分词器
                Map<String, Object> analyzerParams = new HashMap<>();
                analyzerParams.put("type", "chinese");
                builder.analyzerParams(analyzerParams);
                builder.enableAnalyzer(true);
                //开启关键字搜索
                builder.enableMatch(true);
            }
            addFieldReqs.add(builder.build());
        });
        //创建字段请求
        CreateCollectionReq.CollectionSchema collectionSchema = MilvusClientV2.CreateSchema();
        addFieldReqs.forEach(collectionSchema::addField);
        //创建全文检索bm25 函数 默认将text 字段进行稀疏向量化
        collectionSchema.addFunction(CreateCollectionReq
                .Function
                .builder()
                .functionType(FunctionType.BM25)
                .name("text_bm25_emb")
                .inputFieldNames(Collections.singletonList("text"))
                .outputFieldNames(Collections.singletonList("sparse_vector"))
                .build());
        return collectionSchema;
    }


    /**
     * 创建集合中字段索引
     *
     * @return 索引参数
     */
    public static List<IndexParam> createIndex() {
        List<IndexParam> indexParams = new ArrayList<>();
        Map<String, Object> params = new HashMap<>();
        //稀疏向量索引定义
        params.put("inverted_index_algo", "DAAT_MAXSCORE");
        params.put("bm25_k1", 1.2);
        params.put("bm25_b", 0.20);
        IndexParam sparseBuild = IndexParam.builder()
                .fieldName("sparse_vector")
                .indexName("sparse_vector_index")
                .indexType(IndexParam.IndexType.SPARSE_INVERTED_INDEX)
                .metricType(IndexParam.MetricType.BM25)
                .extraParams(params)
                .build();
        indexParams.add(sparseBuild);
        //密集向量索引定义
        IndexParam vectorBuild = IndexParam.builder()
                .indexName("vector_index")
                .fieldName("vector")
                .indexType(IndexParam.IndexType.IVF_FLAT)
                .metricType(IndexParam.MetricType.IP)
                .build();
        indexParams.add(vectorBuild);

        //元数据json 字段定义索引
        //milvus 2.5 不支持自动索引，2.6 可以设置自动索引
       /* IndexParam metadataBuild = IndexParam.builder()
                .indexName("metadata_index")
                .fieldName("metadata")
                .indexType(IndexParam.IndexType.AUTOINDEX)
                .build();
        indexParams.add(metadataBuild);*/
        return indexParams;
    }

    /**
     * 插入数据
     *
     * @param collectionName 集合名
     * @param dataList       数据
     * @return 结果
     */
    public static boolean insertData(String collectionName, String partitionName, List<JsonObject> dataList) {
        InsertResp insert = getMilvusClientV2().insert(InsertReq.builder()
                .data(dataList)
                .collectionName(collectionName)
                .partitionName(partitionName)
                .build());
        return insert.getInsertCnt() > 0;
    }

    /**
     * 通过id 修改向量
     *
     * @param collectionName 集合名称
     * @param partitionName  分区名称
     * @param data           修改数据
     * @return 结果
     */
    public static boolean updateData(String collectionName, String partitionName, List<JsonObject> data) {
        UpsertResp upsert = getMilvusClientV2().upsert(UpsertReq.builder()
                .data(data)
                .collectionName(collectionName)
                .partitionName(partitionName)
                .build()
        );
        return upsert.getUpsertCnt() > 0;
    }

    /**
     * 通过id 删除分区中的实体
     *
     * @param collectionName 集合名称
     * @param documentId     实体id
     * @return boolean
     */
    public static boolean deleteData(String collectionName, List<Object> documentId) {
        DeleteResp delete = getMilvusClientV2().delete(DeleteReq.builder()
                .collectionName(collectionName)
                .ids(documentId)
                .build());
        return delete.getDeleteCnt() > 0;
    }

    /**
     * 将常用的向量检索封装成工具类方便使用(稀疏向量检索 全文检索)
     */
    public static SearchResp sparseSearchResp(String collectionName,
                                              List<String> partitionName,
                                              String query,
                                              int topK,
                                              List<String> outputFields) {
        Map<String, Object> searchParams = new HashMap<>();
        //搜索过程中要忽略的低重要性词的比例。
        searchParams.put("drop_ratio_search", 0.1);
        return getMilvusClientV2().search(SearchReq.builder()
                .collectionName(collectionName)
                .partitionNames(partitionName)
                .data(Collections.singletonList(new EmbeddedText(query)))
                .annsField("sparse_vector")
                .topK(topK)
                .searchParams(searchParams)
                .outputFields(outputFields)
                .build());
    }

    /**
     * 将常用的向量检索封装成工具类方便使用（密集向量检索）
     */
    public static SearchResp searchResp(String collectionName,
                                        List<Long> elementIds,
                                        List<String> partitionName,
                                        List<Float> vector,
                                        int topK,
                                        List<String> outputFields) {
        // 动态构建过滤条件
        StringBuilder filterBuilder = new StringBuilder();
        if (!elementIds.isEmpty()) {
            filterBuilder.append("id in [");
            filterBuilder.append(CollectionUtil.join(elementIds, ","));
            filterBuilder.append("]");
        }
        String filter = filterBuilder.toString();
        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .outputFields(outputFields)
                .topK(topK)
                .data(Collections.singletonList(new FloatVec(vector)))
                .annsField("vector")
                .filter(filter)
                .searchParams(Map.of("nprobe", 10, "metric_type", "IP"))
                .build();
        if (!partitionName.isEmpty()) {
            searchReq.setPartitionNames(partitionName);
        }
        return getMilvusClientV2().search(searchReq);
    }

    /**
     * 混合搜索
     *
     * @param searchRequests   搜索请求
     * @param reranker         重排序策略
     * @param partitionNames   分区名称
     * @param collectionName   搜索集合名称
     * @param topK             取前多少条数据
     * @param consistencyLevel 一致性级别  默认给ConsistencyLevel.BOUNDED
     *                         <html>
     *                         使用 WeightedRanker 策略时，需要在WeightedRanker
     *                         函数中输入权重值。混合搜索中的基本 ANN 搜索次数与需要输入的值的次数相对应。
     *                         输入值的范围应为 [0,1]，数值越接近 1 表示重要性越高。
     *                         BaseRanker WeightedRanker= new  WeightedRanker(Arrays.asList(0.8f, 0.3f));
     *                         <p>
     *                         使用 RRFRanker 策略时，需要将参数值k 输入 RRFRanker。k 的默认值为 60。
     *                         该参数有助于确定如何组合来自不同 ANN 搜索的排名，目的是平衡和混合所有搜索的重要性。
     *                         BaseRanker RRFRanker = new RRFRanker(100);
     *                         </html>
     */
    public SearchResp annSearch(List<AnnSearchReq> searchRequests,
                                List<String> outFiles,
                                List<String> partitionNames,
                                BaseRanker reranker,
                                int topK,
                                String collectionName,
                                ConsistencyLevel consistencyLevel) {

        HybridSearchReq hybridSearchReq = HybridSearchReq.builder()
                .collectionName(collectionName)
                .partitionNames(partitionNames)
                .searchRequests(searchRequests)
                .outFields(outFiles)
                .ranker(reranker)
                .topK(topK)
                .consistencyLevel(consistencyLevel)
                .build();

        return getMilvusClientV2().hybridSearch(hybridSearchReq);
    }


    /**
     * 混合检索示例（结合全文检索和密集向量 来提高召回精度）
     */
    public SearchResp exampleSearch(List<String> partitionNames,
                                    List<String> outFiled,
                                    int topK,
                                    String collectionName,
                                    Map<Long, Float> sparseVector,
                                    List<Float> vector) {
        TreeMap<Long, Float> longFloatTreeMap = new TreeMap<>(sparseVector);

        List<BaseVector> queryDenseVectors = Collections.singletonList(new FloatVec(vector));
        List<BaseVector> querySparseVectors = Collections.singletonList(new SparseFloatVec(longFloatTreeMap));
        List<AnnSearchReq> searchRequests = new ArrayList<>();
        searchRequests.add(AnnSearchReq.builder()
                .vectorFieldName("content_vector_deep")
                .vectors(queryDenseVectors)
                .metricType(IndexParam.MetricType.IP)
                .params("{\"nprobe\": 10}")
                .topK(10)
                .build());
        searchRequests.add(AnnSearchReq.builder()
                .vectorFieldName("content_sparse_vector_deep")
                .vectors(querySparseVectors)
                .metricType(IndexParam.MetricType.IP)
                .params("{}")
                .topK(10)
                .build());
        return annSearch(searchRequests, outFiled, partitionNames, new WeightedRanker(Arrays.asList(0.6f, 0.4f)), topK, collectionName, ConsistencyLevel.BOUNDED);
    }

    /**
     * 向量检索灵活使用
     */
    public static SearchResp search(SearchReq searchReq) {
        return getMilvusClientV2().search(searchReq);
    }

}
